{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3bd19d2c",
   "metadata": {},
   "source": [
    "## Evaluation of showers submitted to the Fast Calorimeter Challenge 2022"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e7163f",
   "metadata": {},
   "source": [
    "This is an interactive version of the file ```evaluate.py```. It can also be run directly using:\n",
    "\n",
    "```\n",
    "python evaluate.py -i INPUT_FILE -r REFERENCE_FILE -m MODE -d DATASET --output_dir OUTPUT_DIR --source_dir SOURCE_DIR\n",
    "```\n",
    "\n",
    "where the arguments are:\n",
    "\n",
    "- ```INPUT_FILE``` is the .hdf5 file that contains the showers to be evaluated.\n",
    "\n",
    "- ```REFERENCE_FILE``` is either the .hdf5 file that contains the showers the input is compared to, or the .pkl file that is created when the code is run for the first time. The latter contains all relevant high-level features and using it results in a faster runtime.\n",
    "\n",
    "- ```MODE``` is one of [all, avg, avg-E, hist, hist-p, hist-chi], and defaults to 'all'. 'avg' plots the average shower of all provided events; 'avg-E' plots the average shower in smaller energy ranges; 'hist' plots histograms of high-level features and saves the separation power, a measure of difference between the histogram of the provided file and the histogram of a reference in SOURCE_DIR, into a file; 'hist-p' only plots the histograms; 'hist-chi' only saves the separation power; and 'all' does all of the above.\n",
    "\n",
    "- ```DATASET``` is the name of the dataset that should be evaluated. Must be one of [1-photons, 1-pions, 2, 3].\n",
    "\n",
    "- ```OUTPUT_DIR``` is the folder in which the plots and other files will be stored. It defaults to 'evaluation_results/'.\n",
    "\n",
    "- ```SOURCE_DIR``` is the folder in which the reference .pkl will be stored. In the future, it will also be where the .hdf5 files for the classifier are saved\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76a6eefc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import evaluate\n",
    "import argparse\n",
    "import h5py\n",
    "import numpy as np\n",
    "\n",
    "import HighLevelFeatures as HLF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bfbf7f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify to your needs:\n",
    "\n",
    "INPUT_FILE = '../data/eval-1-photon.hdf5' # REPLACE THIS WITH YOUR GENERATED EVENTS\n",
    "REFERENCE_FILE = '../data/dataset_1_photons_2.hdf5' # These are the GEANT evaluation events that are provided on zenodo\n",
    "#REFERENCE_FILE = 'source/dataset_2_2.pkl' # This is computed in the first run of the notebook. It can be used instead of the .hdf5 to save time in subsequent runs.\n",
    "MODE = 'all' # not really needed here because the nb is interactive\n",
    "DATASET = '1-photons'\n",
    "OUTPUT_DIR = 'evaluation_results/'\n",
    "SOURCE_DIR = '../data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33985768",
   "metadata": {},
   "outputs": [],
   "source": [
    "# emulating the argument parser of evaluate.py\n",
    "parser_replacement = {\n",
    "    'input_file': INPUT_FILE, 'reference_file': REFERENCE_FILE, 'mode': MODE, 'dataset': DATASET, \n",
    "    'output_dir': OUTPUT_DIR, 'source_dir': SOURCE_DIR, }\n",
    "args = argparse.Namespace(**parser_replacement)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54fc6ea0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reading in source file\n",
    "source_file = h5py.File(args.input_file, 'r')\n",
    "\n",
    "# checking if it has correct shape\n",
    "evaluate.check_file(source_file, args)\n",
    "\n",
    "# preparing output directory\n",
    "if not os.path.isdir(args.output_dir):\n",
    "    os.makedirs(args.output_dir)\n",
    "\n",
    "# preparing source directory\n",
    "if not os.path.isdir(args.source_dir):\n",
    "    os.makedirs(args.source_dir)\n",
    "\n",
    "# extracting showers and energies from source file\n",
    "shower, energy = evaluate.extract_shower_and_energy(source_file, which='input')\n",
    "\n",
    "# creating helper class for high-level features\n",
    "particle = {'1-photons': 'photon', '1-pions': 'pion',\n",
    "            '2': 'electron', '3': 'electron'}[args.dataset]\n",
    "hlf = HLF.HighLevelFeatures(particle, filename='binning_dataset_{}.xml'.format(args.dataset.replace('-', '_')))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e12f60e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reading in reference\n",
    "if os.path.splitext(args.reference_file)[1] == '.hdf5':\n",
    "    print(\"using .hdf5 reference\")\n",
    "    reference_file = h5py.File(args.reference_file, 'r')\n",
    "    evaluate.check_file(reference_file, args, which='reference')\n",
    "    reference_hlf = HLF.HighLevelFeatures(particle, filename='binning_dataset_{}.xml'.format(\n",
    "        args.dataset.replace('-', '_')))\n",
    "    reference_shower, reference_energy = evaluate.extract_shower_and_energy(reference_file,  which='reference')\n",
    "    reference_hlf.Einc = reference_energy\n",
    "    evaluate.save_reference(reference_hlf, args.reference_file)\n",
    "\n",
    "elif os.path.splitext(args.reference_file)[1] == '.pkl':\n",
    "    print(\"using .pkl file for reference\")\n",
    "    reference_hlf = evaluate.load_reference(args.reference_file)\n",
    "else:\n",
    "    raise ValueError(\"reference_file must be .hdf5 or .pkl!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dad888f",
   "metadata": {},
   "source": [
    "### The cells below correspond to different evaluation MODEs and can be run independent of each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d7c06f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation mode 'avg': average of given showers\n",
    "print(\"Plotting average shower...\")\n",
    "_ = hlf.DrawAverageShower(shower, filename=os.path.join(args.output_dir, \n",
    "                                                        'average_shower_dataset_{}.png'.format(args.dataset)),\n",
    "                                  title=\"Shower average\")\n",
    "if hasattr(reference_hlf, 'avg_shower'):\n",
    "    pass\n",
    "else:\n",
    "    reference_hlf.avg_shower = reference_shower.mean(axis=0, keepdims=True)\n",
    "    evaluate.save_reference(reference_hlf, args.reference_file)\n",
    "_ = hlf.DrawAverageShower(reference_hlf.avg_shower, \n",
    "                          filename=os.path.join(args.output_dir, 'reference_average_shower_dataset_{}.png'.format(\n",
    "                                          args.dataset)),\n",
    "                          title=\"Shower average reference dataset\")\n",
    "print(\"Plotting average shower: DONE.\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e48a46f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# evaluation mode 'avg-E': average showers at different energy ranges\n",
    "print(\"Plotting average showers for different energies ...\")\n",
    "if '1' in args.dataset:\n",
    "    target_energies = 2**np.linspace(8, 23, 16)\n",
    "    plot_title = ['shower average at E = {} MeV'.format(int(en)) for en in target_energies]\n",
    "else:\n",
    "    target_energies = 10**np.linspace(3, 6, 4)\n",
    "    plot_title = []\n",
    "    for i in range(3, 7):\n",
    "        plot_title.append('shower average for E in [{}, {}] MeV'.format(10**i, 10**(i+1)))\n",
    "for i in range(len(target_energies)-1):\n",
    "    filename = 'average_shower_dataset_{}_E_{}.png'.format(args.dataset,\n",
    "                                                                   target_energies[i])\n",
    "    which_showers = ((energy >= target_energies[i]) & (energy < target_energies[i+1])).squeeze()\n",
    "    _ = hlf.DrawAverageShower(shower[which_showers],\n",
    "                              filename=os.path.join(args.output_dir, filename),\n",
    "                              title=plot_title[i])\n",
    "    if hasattr(reference_hlf, 'avg_shower_E'):\n",
    "        pass\n",
    "    else:\n",
    "        reference_hlf.avg_shower_E = {}\n",
    "    if target_energies[i] in reference_hlf.avg_shower_E:\n",
    "        pass\n",
    "    else:\n",
    "        which_showers = ((reference_hlf.Einc >= target_energies[i]) & (reference_hlf.Einc < target_energies[i+1])).squeeze()\n",
    "        reference_hlf.avg_shower_E[target_energies[i]] = reference_shower[which_showers].mean(axis=0, keepdims=True)\n",
    "        evaluate.save_reference(reference_hlf, args.reference_file)\n",
    "\n",
    "        _ = hlf.DrawAverageShower(reference_hlf.avg_shower_E[target_energies[i]],\n",
    "                                  filename=os.path.join(args.output_dir, 'reference_'+filename),\n",
    "                                  title='reference '+plot_title[i])\n",
    "\n",
    "print(\"Plotting average shower for different energies: DONE.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b14c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation mode 'hist': plotting histograms of high-level features and printing/saving the sepration power\n",
    "# (equivalent to running hist-p for plotting and hist-chi for the separation power)\n",
    "print(\"Calculating high-level features for histograms ...\")\n",
    "hlf.CalculateFeatures(shower)\n",
    "hlf.Einc = energy\n",
    "\n",
    "print(\"Calculating high-level features for histograms: DONE.\\n\")\n",
    "if reference_hlf.E_tot is None:\n",
    "    reference_hlf.CalculateFeatures(reference_shower)\n",
    "    evaluate.save_reference(reference_hlf, args.reference_file)\n",
    "print(\"Calculating high-level features for histograms: DONE.\\n\")\n",
    "\n",
    "with open(os.path.join(args.output_dir, 'histogram_chi2_{}.txt'.format(args.dataset)), 'w') as f:\n",
    "    f.write('List of chi2 of the plotted histograms, see eq. 15 of 2009.03796 for its definition.\\n')\n",
    "print(\"Plotting histograms ...\")\n",
    "evaluate.plot_histograms(hlf, reference_hlf, args)\n",
    "print(\"Plotting histograms: DONE. \\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a65fa193",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
